﻿using System.Collections;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Runtime.Serialization;
using UnityEngine;
using Utility;
using System;
using System.Linq;

namespace Utility.DataStructure
{
	public class SetDictionary<TKey, TValue> : Dictionary<TKey, HashSet<TValue>>
	{
		public Func<TValue, TKey> keySelector = key => throw new Exception ("KeySelector is null");

		public SetDictionary ()
		{
			AllValues = from set in Values
						from value in set
						select value;
		}

		public SetDictionary (Func<TValue, TKey> keySelector) : this ()
		{
			this.keySelector = keySelector;
		}

		public IEnumerable<TValue> AllValues { get; private set; }

		public HashSet<TValue> GetSet (TKey key, bool createNew = true)
		{
			if (!TryGetValue (key, out var set))
			{
				if (createNew)
				{
					set = new HashSet<TValue> ();
					this[key] = set;
				}
			}
			return set;
		}

		public void ClearSets ()
		{
			foreach (var set in Values)
			{
				set.Clear ();
			}
		}

		public void Add (TKey key, TValue value)
		{
			GetSet (key).Add (value);
		}

		public void Remove (TKey key, TValue value)
		{
			GetSet (key, false).Remove (value);
		}

		public void Clear (TKey key)
		{
			GetSet (key, false).Clear ();
		}

		public bool Contains (TKey key, TValue value)
		{
			return GetSet (key, false)?.Contains (value) ?? false;
		}

		public int GetCount (TKey key)
		{
			return GetSet (key)?.Count ?? 0;
		}


		public void Add (TValue value)
		{
			Add (keySelector (value), value);
		}

		public void Remove (TValue value)
		{
			Remove (keySelector (value), value);
		}

		public bool Contains (TValue value)
		{
			return Contains (keySelector (value), value);
		}

		public ReadOnlySetDictionary<TKey, TValue> ToReadOnly () => new ReadOnlySetDictionary<TKey, TValue> (this);
	}

	public class ReadOnlySetDictionary<TKey, TValue> : IReadOnlyDictionary<TKey, ReadOnlySet<TValue>>
	{
		SetDictionary<TKey, TValue> dict;
		Dictionary<TKey, ReadOnlySet<TValue>> keySetDict = new Dictionary<TKey, ReadOnlySet<TValue>> ();
		Dictionary<ReadOnlySet<TValue>, object> setSetDict = new Dictionary<ReadOnlySet<TValue>, object> ();


		public ReadOnlySet<TValue> GetReadonlySet (TKey key, bool createNew = true)
		{
			dict.TryGetValue (key, out var set);
			keySetDict.TryGetValue (key, out var roSet);

			if (set == null)
			{
				if (roSet != null)
				{
					keySetDict.Remove (key);
					setSetDict.Remove (roSet);
				}
			}
			else
			{
				if (roSet != null && setSetDict[roSet] != set)
				{
					keySetDict.Remove (key);
					setSetDict.Remove (roSet);
					roSet = null;
				}
				if (roSet == null && createNew)
				{
					roSet = new ReadOnlySet<TValue> (set);
					keySetDict[key] = roSet;
					setSetDict[roSet] = set;
				}
			}
			return roSet ?? ReadOnlySet<TValue>.Empty;
		}

		public IEnumerable<TKey> Keys => dict.Keys;

		public IEnumerable<ReadOnlySet<TValue>> Values { get; }

		public int Count => dict.Count;

		public ReadOnlySet<TValue> this[TKey key] => GetReadonlySet (key, false);

		public ReadOnlySetDictionary (SetDictionary<TKey, TValue> dict)
		{
			this.dict = dict;
			Values = from key in dict.Keys select GetReadonlySet (key);
		}



		public bool Contains (TKey key, TValue value)
		{
			return dict.Contains (key, value);
		}

		public int GetCount (TKey key)
		{
			return dict.GetCount (key);
		}

		public bool Contains (TValue value)
		{
			return dict.Contains (value);
		}

		public bool ContainsKey (TKey key)
		{
			return dict.ContainsKey (key);
		}

		public bool TryGetValue (TKey key, out ReadOnlySet<TValue> value)
		{
			value = null;
			if (dict.ContainsKey (key))
			{
				value = GetReadonlySet (key);
				return true;
			}
			return false;
		}

		public IEnumerator<KeyValuePair<TKey, ReadOnlySet<TValue>>> GetEnumerator ()
		{
			foreach (var key in Keys)
			{
				yield return new KeyValuePair<TKey, ReadOnlySet<TValue>> (key, GetReadonlySet (key));
			}
		}

		IEnumerator IEnumerable.GetEnumerator ()
		{
			return GetEnumerator ();
		}
	}
}